#include "stdafx.h"
#include "RayTraceSystem.h"
#include "Scene.h"
#include "Material.h"
#include <utility>
#include "Windows/GxDebug.h"
#include "Basic/StringBuilder.h"
#include "GpuPrimitives.h"

using namespace GxLibBasic;
using namespace GxWinSystem;
namespace RayTracePro
{
	
	const PlatformVendor FirstPreferredPlatform = PlatformVendor::NVidia;
	const PlatformVendor SecondPreferredPlatform = PlatformVendor::AMD;

	int GpuDevice::MaxGlobalWorkItems;
	bool GpuDevice::initialized = false;
	int GpuDevice::scan_WorkGroupSize;
	cl::Kernel GpuDevice::intScanKernel;
	cl::Context GpuDevice::context;
	cl::Program GpuDevice::program;
	cl::CommandQueue GpuDevice::queue;
	cl::Device GpuDevice::device;
	RefPtr<GpuScan> gpuScan;
	__int64 EvalSourceChecksum(String src)
	{
		__int64 rs = 0;
		const char * mbChar = src.ToMBString();
		int len = strlen(mbChar);
		for (int i = 0; i<=len/8; i++)
		{
			__int64 word = 0;
			for (int j = 0; j<8; j++)
			{
				word<<=8;
				if (i*8+j<len)
					word += mbChar[i*8+j];
			}
			rs ^= word;
		}
		return rs;
	}

	PlatformVendor GetPlatformVendor(cl::Platform platform)
	{
		auto vendor = String(platform.getInfo<CL_PLATFORM_VENDOR>(0).c_str());
		if (vendor.Find(L"NVIDIA") != -1)
			return PlatformVendor::NVidia;
		else if (vendor.Find(L"Advanced Micro Devices") != -1)
			return PlatformVendor::AMD;
		else if (vendor.Find(L"Intel") != -1)
			return PlatformVendor::Intel;
		return PlatformVendor::Unkown;
	}

	String GpuDevice::GenerateKernelSources(PlatformVendor vendor, cl_device_type deviceType)
	{
		StringBuilder sb(65536);
		switch (vendor)
		{
		case RayTracePro::PlatformVendor::Unkown:
			sb.WriteString(L"#define OCL_PLATFORM_UNKNOW\n");
			break;
		case RayTracePro::PlatformVendor::NVidia:
			sb.WriteString(L"#define OCL_PLATFORM_NVIDIA\n");
			break;
		case RayTracePro::PlatformVendor::AMD:
			sb.WriteString(L"#define OCL_PLATFORM_AMD\n");
			break;
		case RayTracePro::PlatformVendor::Intel:
			sb.WriteString(L"#define OCL_PLATFORM_INTEL\n");
			break;
		default:
			sb.WriteString(L"#define OCL_PLATFORM_UNKNOW\n");
			break;
		}
		if (deviceType == CL_DEVICE_TYPE_GPU)
			sb.WriteString(L"#define OCL_DEVICE_GPU\n");
		else
			sb.WriteString(L"#define OCL_DEVICE_CPU\n");
		sb.WriteString(ReadTextFile(L"GpuDataStructures.cl"));
		sb.WriteChar(L'\n');
		sb.WriteString(ReadTextFile(L"GpuDataStructures_IO.cl"));
		sb.WriteChar(L'\n');
		sb.WriteString(ReadTextFile(L"kernels.cl"));
		sb.WriteChar(L'\n');
		sb.WriteString(ReadTextFile(L"IntScan.cl"));
		sb.WriteChar(L'\n');
		sb.WriteString(ReadTextFile(L"Scan.cl"));
		return sb.MoveToString();
	}

	void GpuDevice::Initialize()
	{
		PlatformVendor chosenVendor = PlatformVendor::Unkown;
		if (!GpuDevice::initialized)
		{
			printf("Initializing GPU Device...\n");
			unsigned int n;
			cl_platform_id platformIds[10];
			clGetPlatformIDs(10, platformIds, &n);
			cl::vector<cl::Platform> platforms;
			cl::Platform::get(&platforms);
			cl::vector<cl::Device> devices;
			cl::Platform chosenPlatform;
			bool preferredPlatformFound = false;

			for (int i = 0; i<(int)platforms.size(); i++)
			{
				auto name = platforms[i].getInfo<CL_PLATFORM_NAME>(0);
				auto vendor = GetPlatformVendor(platforms[i]);
				auto vendorStr = platforms[i].getInfo<CL_PLATFORM_VENDOR>(0);
				printf("  Platform %d: %s  Vendor: %s\n", i, name.c_str(), vendorStr.c_str());
				if (vendor == FirstPreferredPlatform)
				{
					chosenPlatform = platforms[i];
					chosenVendor = vendor;
					preferredPlatformFound = true;
				}
				else if (vendor == SecondPreferredPlatform && !preferredPlatformFound)
				{
					chosenVendor = vendor;
					chosenPlatform = platforms[i];
					preferredPlatformFound = true;
				}
			}
			if (!preferredPlatformFound)
			{
				chosenVendor = GetPlatformVendor(platforms[0]);
				chosenPlatform = platforms[0];
			}
			printf("Chosen Platform: %s\n", chosenPlatform.getInfo<CL_PLATFORM_NAME>(0).c_str());
			chosenPlatform.getDevices(CL_DEVICE_TYPE_ALL, &devices);
			for (int i = 0; i < (int)devices.size(); i++)
			{
				char * typeStr = "Unkown";
				auto devType = devices[i].getInfo<CL_DEVICE_TYPE>();
					
				switch (devType)
				{
				case CL_DEVICE_TYPE_CPU:
					typeStr = "CPU";
					break;
				case CL_DEVICE_TYPE_GPU:
					typeStr = "GPU";
					break;
				default:
					break;
				}
				printf("  Device %d: %s (%s)\n", i, devices[i].getInfo<CL_DEVICE_NAME>().c_str(), typeStr);
			}
			cl::vector<cl::Device> chosenDevice;
			chosenDevice.push_back(devices[0]);
			device = devices[0];
			MaxGlobalWorkItems = 1<<22;
			context = cl::Context(chosenDevice);
			queue = cl::CommandQueue(context, devices[0]);
			bool needRebuild = true;
			auto source = GenerateKernelSources(chosenVendor, device.getInfo<CL_DEVICE_TYPE>());
			__int64 sourceCheckSum = EvalSourceChecksum(source);
			if (FileExists(L"kernels.bin"))
			{
				printf("Loading kernels...");
				FileStream fs(L"kernels.bin", gxfmRead);
				String deviceName;
				deviceName.LoadFromStream(&fs);
				if (deviceName == devices[0].getInfo<CL_DEVICE_NAME>().c_str())
				{
					int srcLen = 0;
					fs.ReadBuffer(&srcLen, 1);
					__int64 checkSum;
					fs.ReadBuffer(&checkSum, 1);
					if (srcLen == source.GetLength() && checkSum == sourceCheckSum)
					{
						try
						{
							Array<char, true> binary;
							int binLen;
							fs.ReadBuffer(&binLen, 1);
							binary.SetSize(binLen);
							fs.ReadBuffer(binary.Buffer(), 1, binLen);
							cl::Program::Binaries bins;
							bins.push_back(std::make_pair((void*)binary.Buffer(), binLen));
							program = cl::Program(context, chosenDevice, bins);
							
							needRebuild = false;
						}
						catch(...)
						{
							needRebuild = true;
						}
					}
				}
			}
			if (needRebuild)
			{
				// rebuild program from source
				cl::Program::Sources sources(1, std::make_pair(source.ToMBString(), source.GetLength()));
				program = cl::Program(context, sources);	
			}
			try
			{
				printf("  Building kernels...\n");
				char * options;
				if (chosenVendor == PlatformVendor::NVidia)
					options = "-cl-mad-enable -cl-nv-maxrregcount=50";
				else
				{
					options = "-cl-mad-enable";
				}
				program.build(chosenDevice, options);
				printf("Device ready.\n");
				if (needRebuild)
				{
					try
					{
						FileStream fs(L"kernels.bin", gxfmWrite);
						String deviceName = chosenDevice[0].getInfo<CL_DEVICE_NAME>().c_str();
						deviceName.SaveToStream(&fs);

						int srcLen = source.GetLength();
						fs.WriteBuffer(&srcLen, 1);
						fs.WriteBuffer(&sourceCheckSum, 1);
						auto binLengthes = program.getInfo<CL_PROGRAM_BINARY_SIZES>();
						Array<Array<char, true>> binDatas;
						Array<char*, true> pointers;
						pointers.SetSize(binLengthes.size());
						binDatas.SetSize(binLengthes.size());
						for (int i = 0; i<(int)binLengthes.size(); i++)
						{
							binDatas[i].SetSize(binLengthes[i]);
							pointers[i] = binDatas[i].Buffer();
						}
						
						int err = clGetProgramInfo(program(), CL_PROGRAM_BINARIES, binLengthes.size()*sizeof(char*), pointers.Buffer(), 0);
						fs.WriteBuffer(&binLengthes[0], 1);
						fs.WriteBuffer(binDatas[0].Buffer(), binLengthes[0]);
					}
					catch(...)
					{
						printf("  Failed to save compiled kernel binary.\n");
					}
				}
			}
			catch(cl::Error)
			{
				Debug::WriteLine(L"OpenCL Build Error!");
			}
			auto errStr = program.getBuildInfo<CL_PROGRAM_BUILD_LOG>(devices[0], 0);
			WriteTextFile(L"D:\\cl_build.txt", errStr.c_str());
			Debug::WriteLine(errStr.c_str());
		}
		GpuDevice::initialized = true;
		intScanKernel = CreateKernel("kernel__scan_block_anylength");
		scan_WorkGroupSize = intScanKernel.getWorkGroupInfo<CL_KERNEL_WORK_GROUP_SIZE>(device);
		gpuScan = new GpuScan(context(), queue());
	}

	void GpuDevice::Finalize()
	{
		GpuDevice::initialized = false;
		program = cl::Program();
		queue = cl::CommandQueue();
		context = cl::Context();
		gpuScan = 0;
	}

	void GpuDevice::Finish()
	{
		queue.finish();
	}

	void GpuDevice::Barrier()
	{
		queue.enqueueBarrier();
	}

	cl::Kernel GpuDevice::CreateKernel(char * kernelName)
	{
		return cl::Kernel(program, kernelName);
	}

	cl::KernelFunctor GpuDevice::InvokeKernel(cl::Kernel & kernel, cl::NDRange global, cl::NDRange local)
	{
		return kernel.bind(queue, global, local);
	}

	BatchedKernelFunctor GpuDevice::InvokeKernel(cl::Kernel & kernel, int globalSize, int localSize, int batchSize)
	{
		if (batchSize == 0)
			batchSize = MaxGlobalWorkItems;
		localSize = min(localSize, (int)kernel.getWorkGroupInfo<CL_KERNEL_WORK_GROUP_SIZE>(device));
		batchSize = batchSize/localSize * localSize;
		if (batchSize == 0)
			throw "Batch size must be evenly divisible by block size!";
		return BatchedKernelFunctor(queue, kernel, globalSize, localSize, batchSize);
	}

	void GpuDevice::AllocScanList(DeviceList<int> & list, int size)
	{
		//list.SetCapacity(gpuScan->GetScanArrayAllocationSize(size));
		list.SetSize(gpuScan->GetScanArrayAllocationSize(size));
	}

	void GpuDevice::ExclusiveScan(const cl::Buffer & bufferOut, const cl::Buffer & bufferIn, int size)
	{
		gpuScan->Scan(bufferOut(), bufferIn(), size);
	}

	cl::Event GpuDevice::ExclusiveScan(const cl::Buffer & buffer, int offset, int size)
	{
		cl_int clStatus;
		int blockSize = size / scan_WorkGroupSize;
		int B = blockSize * scan_WorkGroupSize;
		if ((size % scan_WorkGroupSize) > 0) { blockSize++; };
		size_t localWorkSize = {scan_WorkGroupSize};
		size_t globalWorkSize = {Align(size / blockSize, scan_WorkGroupSize)};

		clStatus  = clSetKernelArg(intScanKernel(), 0, scan_WorkGroupSize * sizeof(int), 0);
		clStatus |= clSetKernelArg(intScanKernel(), 1, sizeof(cl_mem), &buffer);
		clStatus |= clSetKernelArg(intScanKernel(), 2, sizeof(int), &B);
		clStatus |= clSetKernelArg(intScanKernel(), 3, sizeof(int), &size);
		clStatus |= clSetKernelArg(intScanKernel(), 4, sizeof(int), &blockSize);
		clStatus |= clSetKernelArg(intScanKernel(), 5, sizeof(int), &offset);
		cl_event evt;
		clStatus |= clEnqueueNDRangeKernel(queue(), intScanKernel(), 1, NULL, &globalWorkSize, &localWorkSize, 0, NULL, &evt);
		if (clStatus != CL_SUCCESS)
		{
			throw cl::Error(clStatus, "IntScan");
		}
		cl::Event rs;// = evt;
		return rs;
	}

	cl::CommandQueue GpuDevice::GetQueue()
	{
		return queue;
	}

	cl::Context GpuDevice::GetContext()
	{
		return context;
	}

	RayTraceSystem::RayTraceSystem()
	{
		GeometryBias = 0.1f;
		AntiAliasSamples = 8;
	}

	Scene * RayTraceSystem::CreateScene()
	{
		return new Scene(this);
	}
	Material * RayTraceSystem::CreateMaterial()
	{
		return new Material();
	}

	void RayTraceSystem::InitRenderThread()
	{
		auto counter = new DebugTimming();
		debugCounters.Add(counter);
		TlsSetValue(0, counter);
	}

	void RayTraceSystem::FinalizeRenderThread()
	{
		auto counter = (const DebugTimming*)TlsGetValue(0);
		int fid = debugCounters.Find(counter);
		if (fid != -1)
			debugCounters.Delete(fid);
	}

	void RayTraceSystem::StartBenchmark(const char * name)
	{
		auto counter = (DebugTimming*)TlsGetValue(0);
		counter->Start(name);
	}

	void RayTraceSystem::StopBenchmark()
	{
		auto counter = (DebugTimming*)TlsGetValue(0);
		counter->End();
	}
}